#include "stream_tools.h"

template<
	unsigned W_IO,	// Width of IO N_IO * BIT
	unsigned W_CH,	// Width of CHannel N_CH * BIT
	unsigned HIGHT,	// Input HIGHT
	unsigned WIDTH,	// Input WIDTH
	unsigned N_BATCH,
	unsigned RAM,
	unsigned S = 1,	// Stride
	unsigned PADVAL = 0x00
>
void stream2window(data_stream<W_IO>& in, data_stream<W_IO>& out)
{
	static_assert(W_CH >= W_IO, "stream2window");
	static_assert(W_CH % W_IO == 0, "stream2window");
	constexpr unsigned FOLD = W_CH / W_IO;
	constexpr unsigned WIDTH_OFFSET = 2;
	const ap_uint<W_IO> PAD_AP = PADVAL;

	assert(in.size() == N_BATCH * HIGHT * WIDTH * FOLD);
	assert(out.empty());

	#pragma HLS DATAFLOW

	ap_uint<2> idx[3] = {0, 1, 2};
	#pragma HLS ARRAY_PARTITION variable=idx complete dim=1
	ap_uint<W_IO> line[3][WIDTH][FOLD];
	if(RAM==0)
	{
		#pragma HLS bind_storage variable=line type=ram_2p impl=lutram
	}
	else if(RAM==1)
	{
		#pragma HLS bind_storage variable=line type=ram_2p impl=bram
	}
	else
	{
		#pragma HLS bind_storage variable=line type=ram_2p impl=uram
	}

	for (unsigned w = 0; w < WIDTH + WIDTH_OFFSET; ++w)
	{
		for (unsigned f = 0; f < FOLD; ++f)
		{
			#pragma HLS PIPELINE II=1
			if (w < WIDTH)
			{
				line[idx[2]][w][f] = in.read();
			}
			else
			{
				line[idx[0]][w - WIDTH][f] = in.read();
			}
		}
	}

	// main loop
	for (unsigned b = 0; b < N_BATCH; b++)
	{
		for (unsigned h = 0; h < HIGHT; h++)
		{
			for (unsigned w = 0; w < WIDTH; w++)
			{
				for (unsigned f = 0; f < FOLD; f++)
				{
					#pragma HLS PIPELINE II=9
					// idx shift
					if (w == 0 && f == 0)
					{
						ap_uint<2> it = idx[0];
						idx[0] = idx[1];
						idx[1] = idx[2];
						idx[2] = it;
					}
					// load data
					if (b < N_BATCH - 1 || h < HIGHT - 2 || (h == HIGHT - 2 && w < WIDTH - WIDTH_OFFSET))
					{
						ap_uint<W_IO> in_buf = in.read();
						if (w < WIDTH - WIDTH_OFFSET)
						{
							line[idx[2]][w + WIDTH_OFFSET][f] = in_buf;
						}
						else
						{
							line[idx[0]][w + WIDTH_OFFSET - WIDTH][f] = in_buf;
						}
					}

					// output, II = 9
					// (h == 0 || h == HIGHT - 1 || w == 0 || w == WIDTH - 1) --> padding
					if( h % S == 0 && w % S == 0 )
					{
						out.write((h != 0 && w != 0) ? 					line[idx[0]][w - 1][f] : PAD_AP);
						out.write((h != 0) ? 							line[idx[0]][w + 0][f] : PAD_AP);
						out.write((h != 0 && w != WIDTH - 1)? 			line[idx[0]][w + 1][f] : PAD_AP);
						out.write((w != 0) ? 							line[idx[1]][w - 1][f] : PAD_AP);
						out.write( 										line[idx[1]][w + 0][f]);
						out.write((w != WIDTH - 1) ? 					line[idx[1]][w + 1][f] : PAD_AP);
						out.write((h != HIGHT - 1 && w != 0) ? 			line[idx[2]][w - 1][f] : PAD_AP);
						out.write((h != HIGHT - 1) ? 					line[idx[2]][w + 0][f] : PAD_AP);
						out.write((h != HIGHT - 1 && w != WIDTH - 1) ?	line[idx[2]][w + 1][f] : PAD_AP);
					}
				}
			}
		}
	}

	assert(in.empty());
	assert(out.size() == N_BATCH * HIGHT * WIDTH * FOLD * 9 / 4);
	return;
}

//stream2window expand K
template<
	unsigned W_IO,	// Width of IO N_IO * BIT
	unsigned W_CH,	// Width of CHannel N_CH * BIT
	unsigned N_K,	// # of Kernel
	unsigned HIGHT,	// Input HIGHT
	unsigned WIDTH,	// Input WIDTH
	unsigned N_BATCH,
	unsigned S = 1,	// Stride
	unsigned PADVAL = 0x00
>
void stream2window_ek(data_stream<W_IO>& in, data_stream<W_IO*N_K>& out)
{
	static_assert(W_CH >= W_IO, "stream2window");
	static_assert(W_CH % W_IO == 0, "stream2window");
	constexpr unsigned FOLD = W_CH / W_IO;
	constexpr unsigned WIDTH_OFFSET = 1;
	const ap_uint<W_IO*N_K> PAD_AP = PADVAL;

	assert(in.size() == N_BATCH * HIGHT * WIDTH * FOLD);
	assert(out.empty());

	ap_uint<2> idx[3] = {0, 1, 2};
	#pragma HLS ARRAY_PARTITION variable=idx complete dim=1
	ap_uint<W_IO*N_K> line[3][WIDTH][FOLD] = {0};
	#pragma HLS bind_storage variable=line type=ram_2p impl=lutram
	ap_uint<W_IO*N_K> inbuf[FOLD] = {0};
	ap_uint<W_IO*N_K> outbuf;

	ap_uint<2> it;
	for (unsigned b = 0; b < N_BATCH; b++)
	{
		for (unsigned f = 0; f < FOLD; f++)
		{
		#pragma HLS PIPELINE II=1
			inbuf[f] = 0;
		}

		for (unsigned f = 0; f < FOLD; f++)
		{
		#pragma HLS PIPELINE II=1
			inbuf[f](SLICE(W_IO, 0)) = in.read();
		}
		for (unsigned w = 0; w < WIDTH; ++w)
		{
			for (unsigned f = 0; f < FOLD; f++)
			{
				#pragma HLS PIPELINE II=1
				if (w < WIDTH - 1)
				{
					inbuf[f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 1));
					inbuf[f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 0));
					inbuf[f](SLICE(W_IO, 0)) = in.read();
					line[idx[2]][w][f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 0));
					line[idx[2]][w][f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 1));
					line[idx[2]][w][f](SLICE(W_IO, 0)) = inbuf[f](SLICE(W_IO, 2));
				}
				else if(w == WIDTH - 1)
				{
					inbuf[f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 1));
					inbuf[f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 0));
					inbuf[f](SLICE(W_IO, 0)) = 0;
					line[idx[2]][w][f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 0));
					line[idx[2]][w][f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 1));
					line[idx[2]][w][f](SLICE(W_IO, 0)) = inbuf[f](SLICE(W_IO, 2));
				}
			}
		}

		// main loop
		for (unsigned h = 0; h < HIGHT - 1; h++)
		{
			for (unsigned f = 0; f < FOLD; f++)
			{
			#pragma HLS PIPELINE II=1
				inbuf[f](SLICE(W_IO, 0)) = in.read();
			}
			for (unsigned w = 0; w < WIDTH; w++)
			{
				for (unsigned f = 0; f < FOLD; f++)
				{
				#pragma HLS PIPELINE II=1
					// idx shift
					if(w == 0 && f == 0)
					{
						it = idx[0];
						idx[0] = idx[1];
						idx[1] = idx[2];
						idx[2] = it;
					}
					if(w == 0)
					{
						inbuf[f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 1));
						inbuf[f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 0));
						inbuf[f](SLICE(W_IO, 0)) = in.read();
						line[idx[2]][w][f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 0));
						line[idx[2]][w][f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 1));
						line[idx[2]][w][f](SLICE(W_IO, 0)) = 0;
					}
					else if (w < WIDTH - 1)
					{
						inbuf[f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 1));
						inbuf[f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 0));
						inbuf[f](SLICE(W_IO, 0)) = in.read();
						line[idx[2]][w][f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 0));
						line[idx[2]][w][f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 1));
						line[idx[2]][w][f](SLICE(W_IO, 0)) = inbuf[f](SLICE(W_IO, 2));
					}
					else if(w == WIDTH - 1)
					{
						inbuf[f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 1));
						inbuf[f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 0));
						inbuf[f](SLICE(W_IO, 0)) = 0;
						line[idx[2]][w][f](SLICE(W_IO, 2)) = inbuf[f](SLICE(W_IO, 0));
						line[idx[2]][w][f](SLICE(W_IO, 1)) = inbuf[f](SLICE(W_IO, 1));
						line[idx[2]][w][f](SLICE(W_IO, 0)) = inbuf[f](SLICE(W_IO, 2));
					}

					// output, II = N_K
					// (h == 0 || h == HIGHT - 1 || w == 0 || w == WIDTH - 1) --> padding
					if( h % S == 0 && w % S == 0 )
					{
						outbuf = (h == 0) ? PAD_AP : line[idx[0]][w][f];
						out.write(outbuf);
						outbuf = line[idx[1]][w][f];
						out.write(outbuf);
						outbuf = (h == HIGHT - 1) ? PAD_AP : line[idx[2]][w][f];
						out.write(outbuf);
					}
				}
			}
		}

		unsigned h = HIGHT - 1;
		for (unsigned w = 0; w < WIDTH; w++)
		{
			for (unsigned f = 0; f < FOLD; f++)
			{
				#pragma HLS PIPELINE II=1
				if( h % S == 0 && w % S == 0 )
				{
					outbuf = line[idx[1]][w][f];
					out.write(outbuf);
					outbuf = line[idx[2]][w][f];
					out.write(outbuf);
					outbuf = PAD_AP;
					out.write(outbuf);
				}
			}
		}
	}

	assert(in.empty());
	assert(out.size() == N_BATCH * HIGHT * WIDTH * FOLD * N_K / S / S);
	return;
}